import datetime

import pandas as pd

from rbt.md import MdEngine

class BnKlineMdEngine(MdEngine):
    def __init__(self, path: str):
        super().__init__()
        self.path = path
    
    def prepare_data(self, sym: str, date: datetime.date):
        # 读取PKL文件
        file_path = f"{self.path}/{sym}/{sym}_kline_1s_{date}.pkl"
        df = pd.read_pickle(file_path)
        
        # 添加必要的列
        df['sym'] = sym
        
        # 设置索引为Timestamp
        df.set_index('close_time', inplace=True)
        
        # 注册原始行情数据
        self._register_raw_md(sym, date, df, recover_mo=False)